import asyncio
import os
import re


from agents import (
    Agent,
    ModelSettings,
    Runner,
    RunResult,
    set_tracing_disabled,
)
from agents.mcp import MCPServerSse
from agents.models.openai_provider import OpenAIProvider


set_tracing_disabled(disabled=True)

# model_id = "qwen3-235b-a22b"
model_id = "qwen3:8b"

extra_body = {}

API_KEY = os.getenv("EXAMPLE_API_KEY")
if model_id == "qwen3-235b-a22b":
    extra_body = {"enable_thinking": False}
    BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"
else:
    API_KEY = "ollama"
    BASE_URL = "http://127.0.0.1:11434/v1"

provider = OpenAIProvider(
    api_key=API_KEY,
    base_url=BASE_URL,
    use_responses=False
    )
model = provider.get_model(model_id)


async def get_restaurants(message: str) -> str:
    global model, extra_body
    
    remote_server = MCPServerSse(
        name="Restaurant Server",
        params={
            "url": "http://localhost:8004/sse",
        },
        cache_tools_list=False
    )

    system_message = "Get a list of restaurants (\"餐馆\" in Chinese) near a location, and the input parameter is a comma-separated latitude and longitude string. The list of restaurants should be plain text, don't use JSON or Markdown format."
    if model_id.startswith("qwen3:"):
        system_message = f"{system_message} /no_think"

    async with remote_server as server:
        # Create an agent that uses the MCP server
        agent = Agent(
            name="Assistant",
            instructions=system_message,
            model=model,
            model_settings=ModelSettings(extra_body=extra_body),
            mcp_servers=[server]
        )

        # Run the agent
        result: RunResult = await Runner.run(agent, message)
        answer = result.final_output
        answer = re.sub(r"<think>.*?</think>\n?", "", answer, flags=re.DOTALL)
        print(answer)
        return answer


if __name__ == "__main__":
    user_message = "我所在位置的经纬度是\"103.91941,33.26922\", 这附近有哪些餐馆？"
    if model_id.startswith("qwen3:"):
        user_message = f"{user_message} /no_think"
        
    asyncio.run(get_restaurants(user_message))
